import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
#数据加载
transform=transforms.Compose([
    transforms.ToTensor(),#转化为张量
    transforms.Normalize((0.5,),(0.5,)),#归一化到[-1,1]
])
#加载MNIST数据集
train_dataset=datasets.MNIST(root="./data",train=True,transform=transform)
test_dataset=datasets.MNIST(root='./data',train=False,transform=transform)

train_loader=DataLoader(dataset=train_dataset,batch_size=64,shuffle=True)
test_loader=DataLoader(dataset=test_dataset,batch_size=64,shuffle=True)

#定义CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN,self).__init__()
#定义输入1通道输出32通道的的卷积层
        self.conv1=nn.Conv2d(1,32,kernel_size=3,stride=1,padding=1)
#定义输入32通道输出64通道的
        self.conv2=nn.Conv2d(32,64,kernel_size=3,stride=1,padding=1)
#定义全链接层
        self.fc1=nn.Linear(64*7*7,128)#输入大小等于特征图大小*通道数
        self.fc2=nn.Linear(128,10)#10个类别
    def forward(self,x):
        x=F.relu(self.conv1(x))#第一层卷积+relu
        x=F.max_pool2d(x,2)#最大池化
        x=F.relu(self.conv2(x))#第二层卷积+relu
        x=F.max_pool2d(x,2)#最大池化
        x=x.view(-1,64*7*7)#展平操作
        x=F.relu(self.fc1(x))#全连接层+Relu
        x=self.fc2(x)#全连接层输出
        return x
model=SimpleCNN()

criterion=nn.CrossEntropyLoss()#多酚类交叉熵损失
optimizer=optim.SGD(model.parameters(),lr=0.001,momentum=0.9)#学习率和动量

num_epochs=5
model.train()#设为训练模式
for epoch in range(num_epochs):
    total_loss=0
    for images,labels in train_loader:
        #向前传播
        outputs=model(images)
        loss=criterion(outputs,labels)
        #向后传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
    print(f'Epoch:[{epoch+1}/{num_epochs}],Loss:{total_loss /len(train_loader):.4f}')

model.eval()
correct=0
total=0
with torch.no_grad():
    for iamges,label in test_loader:
        outputs=model(images)
        _,predicted=torch.max(outputs,1)#预测类型
        total+=labels.size(0)
        correct+=(predicted==labels).sum().item()
accuray=100*correct/total
print(f'Test Accuray:{accuray:.2f}%')


